import numpy as np
try:
    from .typedef import spmOp_t
except:
    from typedef import spmOp_t

RADAR_CYCLE = 0.05
TRK_OBJECT_NUM = 128
START_FRAME = 3
MIN_MOVED_DIST = 5
MAX_EXTRA = 3
MAX_STOP_HOLD = 5
TRK_MIN_SPEED = 0.3
TRK_MAHA_DIST_THLD = 9.0
TRK_INITIATION_MAX_ABS_VREL = 9
STAT = {'terminate':0, 'tentative': 1, 'new': 2, 'confirmed': 3}
_field = {
    'id': 'i',
    'status': 'i',
    'distX': 'f',
    'distY': 'f',
    'distZ': 'f',
    'distXVar': 'f',
    'distYVar': 'f',
    'distZVar': 'f',
    'vRelX': 'f',
    'vRelY': 'f',
    'range': 'f',
    'velocity': 'f',
    'azimuth': 'f',
    'snr': 'f',
    'firstX': 'f',
    'firstY': 'f',
    'sumR': 'f',
    'detect_cnt': 'i',
    'lost_cnt': 'i',
    'pred_cnt': 'i'
}
trk_object_t = np.dtype({'names': tuple(_field.keys()),'formats': tuple(_field.values()),'align': True})

_field = {
    'id': 'i',
    'trkCnt': 'i',
    'distX': 'f',
    'distY': 'f',
    'distZ': 'f',
    'vRelX': 'f',
    'vRelY': 'f',
    'sumR': 'f',
    'sumY': 'f',
    'snr': 'f'
}
trkOp_t = np.dtype({'names': tuple(_field.keys()),'formats': tuple(_field.values()),'align': True})

_field = {
    'x': 'f',
    'y': 'f',
    'vRelX': 'f',
    'vRelY': 'f',
    'range': 'f',
    'velocity': 'f',
    'azimuth': 'f',
    'snr': 'f',
    'distXVar': 'f',
    'distYVar': 'f',
    'classID': 'i',
    'idx': 'i'
}
point_t = np.dtype({'names': tuple(_field.keys()),'formats': tuple(_field.values()),'align': True})

def dbscan(X, eps, min_samples):
    """
    X: 样本数据，形状为 (n_samples, n_features)
    eps: 邻域半径
    min_samples: 核心点的最小样本数
    """
    n_samples = X.shape[0]
    # 初始化标签，-1 表示未访问，0 表示噪声，其他数字表示聚类标签
    labels = np.full(n_samples, -1)

    # 为每个点找到其eps邻域内的点
    def find_neighbors(index, eps):
        distances = np.linalg.norm(X[index] - X, axis=1)
        neighbors = np.where(distances <= eps)[0]
        return neighbors

    # 构建邻接图
    for i in range(n_samples):
        if labels[i] != -1:
            continue  # 如果已经访问过或标记为噪声，跳过
        neighbors = find_neighbors(i, eps)
        if len(neighbors) < min_samples:
            labels[i] = 0  # 标记为噪声
            continue

        # 找到密度可达的所有点
        cluster_label = labels.max() + 1
        labels[i] = cluster_label
        queue = neighbors.tolist()
        while queue:
            current = queue.pop(0)
            if labels[current] != -1:
                continue
            labels[current] = cluster_label
            new_neighbors = find_neighbors(current, eps)
            queue.extend(new_neighbors.tolist())

    return labels

def glueSpmOutput(spmOp: np.ndarray):
    spmOp.dtype = spmOp_t
    point = np.zeros(len(spmOp), dtype=point_t)
    point['range'] = spmOp['range']
    point['velocity'] = spmOp['velocity']
    point['azimuth'] = spmOp['azimuth']
    point['snr'] = spmOp['SNR']
    point['x'] = spmOp['fX']
    point['y'] = spmOp['fY']
    aziRad = np.radians(point['azimuth'])
    point['vRelX'] = point['velocity'] * np.sin(aziRad)
    point['vRelY'] = point['velocity'] * np.cos(aziRad)
    point['distXVar'] = 0.8
    point['distYVar'] = 0.5
    point['classID'] = -1
    return point

defaultTrkCfg = {
    "MATCH_DIST" :      TRK_MAHA_DIST_THLD,
    "MIN_SPEED":        TRK_MIN_SPEED,
    "START_FRAME":      START_FRAME,
    "MIN_MOVED_DIST":   MIN_MOVED_DIST,
}

class Tracker(object):
    def __init__(self):
        t = RADAR_CYCLE
        # Kalman Filter TRACK parameter
        # X, Y, VX, VY
        self.Xe = np.zeros([TRK_OBJECT_NUM, 4, 1], dtype='f')
        self.XeT = np.zeros([TRK_OBJECT_NUM, 4, 1], dtype='f')
        self.Pe = np.zeros([TRK_OBJECT_NUM, 4, 4], dtype='f')
        self.PeT = np.zeros([TRK_OBJECT_NUM, 4, 4], dtype='f')
        self.S = np.zeros([TRK_OBJECT_NUM, 4, 4], dtype='f')
        self.K = np.zeros([TRK_OBJECT_NUM, 4, 4], dtype='f')
        self.F = np.eye(4)
        self.F[0, 2] = t
        self.F[1, 3] = t
        self.Q = np.eye(4)
        self.I = np.eye(4)
        self.P = np.array([[10, 0, 0, 0], [0, 3, 0, 0], [0, 0, 40, 0], [0, 0, 0, 10]])
        self.H = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
        self.R = np.array([[8, 0, 0, 0], [0, 8, 0, 0], [0, 0, 10, 0], [0, 0, 0, 10]])
        self.trace = [[] for _ in range(TRK_OBJECT_NUM)]
        self.obj = np.zeros(TRK_OBJECT_NUM, dtype=trk_object_t)
        self.target = None
        self.cfg = defaultTrkCfg.copy()
        self.objId = 0
        self.TRK_Reset()
        
    def TRK_staticProcess(self, pointCloud: np.ndarray):
        # static class: 0 for microMove, -2 for stability static, -1 for temporary static
        # static target for Grid
        point = glueSpmOutput(pointCloud)
        pointStatic = point[abs(point['velocity'])<0.3]
        pointMove = point[abs(point['velocity'])>0.3]
        return pointMove, pointStatic
    
    def TRK_JPDA(self, track, point):
        vRelThld = 2
        aziThld = 6
        pntSize = len(point)
        tarSize = len(track)
        tarX, tarY, tarV, tarA = track['distX'], track['distY'], track['velocity'], track['azimuth']
        pntX, pntY, pntV, pntA = point['x'], point['y'], point['velocity'], point['azimuth']
        diffX = np.repeat(tarX, pntSize) - np.tile(pntX, tarSize)
        diffY = np.repeat(tarY, pntSize) - np.tile(pntY, tarSize)
        dist = (abs(diffX) + abs(diffY)).reshape([tarSize, pntSize])
        diffA = np.repeat(tarA, pntSize) - np.tile(pntA, tarSize)
        distA = np.degrees(abs(diffA)).reshape([tarSize, pntSize])
        distV = np.repeat(tarV, pntSize) - np.tile(pntV, tarSize)
        distV = abs(distV).reshape([tarSize, pntSize])
        probs = 0.74816 ** (dist)
        probs = np.where(distV > vRelThld, 0, probs)
        probs = np.where(distA > aziThld, 0, probs)
        probs = np.where(dist > TRK_MAHA_DIST_THLD, 0, probs)
        probs = probs.T / (1e-10+np.sum(probs, axis=1))
        return probs.T

    def TRK_Reset(self):
        self.target = None
        self.obj[:] = 0
    
    @staticmethod  
    def TRK_Cluster(pnt):
        minSample = 1
        cluster_dist = 2
        if len(pnt) < minSample:
            return pnt
        cluster2dbscan = np.array([pnt['x'], pnt['y']]).T
        dbscan_result = dbscan(cluster2dbscan, eps=cluster_dist, min_samples=minSample)
        pnt['classID'] = dbscan_result
        cluster = pnt[pnt['classID'] >= 0]  # moving target for tracking
        return cluster
    
    @staticmethod  
    def TRK_Merge(cluster):
        idList = np.unique(cluster["classID"])
        MergeCluster = []
        for i in idList:
            iterCluster = cluster[cluster['classID']==i]
            idx = np.argmax(iterCluster["snr"])
            MergeCluster.append(iterCluster[idx])
        MergeCluster = np.array(MergeCluster, dtype=point_t)
        return MergeCluster

    @staticmethod
    def TRK_Match(target, point, rGate, vGate, aGate):
        tarSize = len(target)
        pntSize = len(point)
        tarR, tarA, tarV = target['range'], target['azimuth'], target['velocity']
        pntR, pntA, pntV = point['range'], point['azimuth'], point['velocity']
        diffR = np.repeat(tarR, pntSize) - np.tile(pntR, tarSize)
        distR = abs(diffR).reshape([tarSize, pntSize])
        diffA = np.repeat(tarA, pntSize) - np.tile(pntA, tarSize)
        distA = np.degrees(abs(diffA)).reshape([tarSize, pntSize])
        diffV = np.repeat(tarV, pntSize) - np.tile(pntV, tarSize)
        distV = abs(diffV).reshape([tarSize, pntSize])
        return np.any(distR < rGate, axis=0) & np.any(distV < vGate, axis=0) & np.any(distA < aGate, axis=0)
    
    def TRK_CreateNewObjects(self, confirmed: np.ndarray):
        stat = self.obj['status']
        clusters = confirmed
        clusters['idx'] = np.arange(len(clusters))
        dummyIdx = np.argwhere(np.isin(stat, STAT['terminate'])).flatten()
        clusterList = np.unique(clusters['classID'])
        max_indices = []
        for classID in clusterList:
            _ = clusters[clusters['classID'] == classID]
            max_indices.append(_['idx'][np.argmax(_['snr'])])
        max_indices = np.array(max_indices)
        valid_dummyIdx = dummyIdx[:len(max_indices)]
        if len(valid_dummyIdx) < len(max_indices):
            print("object list full")
            max_indices = max_indices[:len(dummyIdx)]
        for max_index, idx in zip(max_indices, valid_dummyIdx):
            val = clusters[max_index]
            objIdx = idx
            self.obj[objIdx] = 0
            self.obj['id'][objIdx] = self.objId
            self.objId += 1
            self.obj['status'][objIdx] = STAT['tentative']
            self.Xe[objIdx, 0, 0] = self.XeT[objIdx, 0, 0] = val['x']
            self.Xe[objIdx, 1, 0] = self.XeT[objIdx, 1, 0] = val['y']
            self.Xe[objIdx, 2, 0] = self.XeT[objIdx, 2, 0] = val['vRelX']
            self.Xe[objIdx, 3, 0] = self.XeT[objIdx, 3, 0] = val['vRelY']
            self.Pe[objIdx, :, :] = self.PeT[objIdx, :, :] = self.P
            self.obj['distX'][objIdx] = self.obj['firstX'][objIdx] = val['x']
            self.obj['distY'][objIdx] = self.obj['firstY'][objIdx] = val['y']
            self.obj['vRelX'][objIdx] = val['vRelX']
            self.obj['vRelY'][objIdx] = val['vRelY']
            self.obj['range'][objIdx] = val['range']
            self.obj['velocity'][objIdx] = val['velocity']
            self.obj['azimuth'][objIdx] = val['azimuth']
            self.obj['snr'][objIdx] = val['snr']
            self.obj['distXVar'][objIdx] = val['distXVar']
            self.obj['distYVar'][objIdx] = val['distYVar']
            self.obj['sumR'][objIdx] = 0
            self.trace[objIdx] = []
    
    def TRK_Update(self, measuredIdx, measurements, measured_probs):
        meas_columns = ['x', 'y', 'vRelX', 'vRelY']
        measXe = np.column_stack([measurements[_n] for _n in meas_columns])
        # 预先计算可以重复使用的项
        H_PeT_HT = self.H @ self.PeT @ self.H.T
        I_H = self.I - self.H
        for i_idx, i in enumerate(measuredIdx):
            prob = measured_probs[i_idx, :]
            maxIdx = np.argmax(prob)
            Y = np.sum(measXe.T * prob, axis=1)[:, np.newaxis] - self.H @ self.XeT[i]
            S = H_PeT_HT[i] + self.R
            SI = np.linalg.inv(S)
            K = SI @ self.PeT[i] @ self.H.T  # 注意这里使用了 S 的逆矩阵
            self.Xe[i] = self.XeT[i] + K @ Y
            self.Pe[i] = I_H @ self.PeT[i]
            # 从 measurements 中获取最大概率对应的数据
            dat = measurements[maxIdx]
            # 更新 self.obj ndarray
            self.obj['sumR'][i] += abs(dat['range'] - self.obj['range'][i])
            self.obj['distX'][i] = self.Xe[i,0,0]
            self.obj['distY'][i] = self.Xe[i,1,0]
            self.obj['vRelX'][i] = self.Xe[i,2,0]
            self.obj['vRelY'][i] = self.Xe[i,3,0]
            for _n in ['range', 'velocity', 'azimuth', 'snr']:
                self.obj[_n][i] = dat[_n]
            
    def main_periodic_process(self, pointCloud, t, vehicleInfo):
        ################### Kalman Filter for Moving Target ###################
        self.F[0, 2] = t
        self.F[1, 3] = t
        ################### Static Step ###################
        pointMove, pointStatic = self.TRK_staticProcess(pointCloud)
        cluster = self.TRK_Cluster(pointMove)
        cluster = self.TRK_Merge(cluster)
        ################### object Prediction Step ###################
        stat = self.obj['status']
        _predname = [STAT['confirmed'], STAT['new'], STAT['tentative']]
        predObjIdx = np.argwhere(np.isin(stat, _predname)).flatten()
        for idx in predObjIdx:
            self.Xe[idx,0,0] = self.obj['distX'][idx]
            self.Xe[idx,1,0] = self.obj['distY'][idx]
            self.Xe[idx,2,0] = self.obj['vRelX'][idx]
            self.Xe[idx,3,0] = self.obj['vRelY'][idx]
            self.XeT[idx, :, :] = self.F @ self.Xe[idx, :,:]
            self.PeT[idx, :] = self.F @ self.Pe[idx, :] @ self.F.T + self.Q
        self.obj['distX'][predObjIdx] = self.XeT[predObjIdx, 0, 0]
        self.obj['distY'][predObjIdx] = self.XeT[predObjIdx, 1, 0]
        self.obj['vRelX'][predObjIdx] = self.XeT[predObjIdx, 2, 0]
        self.obj['vRelY'][predObjIdx] = self.XeT[predObjIdx, 3, 0]
        ################### Measurement Step #########################
        measurements = cluster.copy()
        ################### point Association Step ###################
        activeIdx = np.argwhere(~np.isin(stat, STAT['terminate'])).flatten()
        activeObj = self.obj[activeIdx]
        association_probs = self.TRK_JPDA(activeObj, measurements)
        noMatch = np.all(association_probs ==0, axis=1)
        _measSubIdx = np.argwhere(~noMatch).flatten()
        _predSubIdx = np.argwhere(noMatch).flatten()
        measuredIdx = activeIdx[_measSubIdx]
        predictedIdx = activeIdx[_predSubIdx]
        measured_probs = association_probs[_measSubIdx]
        self.obj['detect_cnt'][measuredIdx] += 1
        self.obj['pred_cnt'][measuredIdx] = 0
        _tmp = (self.XeT[predictedIdx,:,0] + self.Xe[predictedIdx,:,0]) / 2
        self.obj['distX'][predictedIdx] = _tmp[:,0]
        self.obj['distY'][predictedIdx] = _tmp[:,1]
        self.obj['vRelX'][predictedIdx] = _tmp[:,2]
        self.obj['vRelY'][predictedIdx] = _tmp[:,3]
        self.obj['lost_cnt'][predictedIdx] += 1
        self.obj['pred_cnt'][predictedIdx] += 1
        ################### object Update Step ###################
        self.TRK_Update(measuredIdx, measurements, measured_probs)
        ################### Terminate Step #################
        termTentativeIdx = np.argwhere((self.obj["status"]==STAT['tentative'])&(self.obj["lost_cnt"]!=0)).flatten()
        termConfirmedIdx = np.argwhere((self.obj["status"]==STAT['confirmed'])&(self.obj["pred_cnt"]>MAX_EXTRA)).flatten()
        self.obj['status'][termTentativeIdx] = STAT['terminate']
        self.obj['status'][termConfirmedIdx] = STAT['terminate']
        ################### Confirmed Step #################
        confirmedIdx = np.argwhere((self.obj["status"]==STAT['tentative'])&(self.obj["detect_cnt"]>=START_FRAME)).flatten()
        self.obj['status'][confirmedIdx] = STAT['confirmed']
        ##################### Start Step ###################
        if len(association_probs) == 0:
            remainPoint = measurements
        else:
            remainPoint = measurements[np.all(association_probs==0, axis=0)]
        if self.target is None:
            self.target = remainPoint
        else:
            rGate = TRK_INITIATION_MAX_ABS_VREL * t
            vGate = 3
            aGate = 5
            bMatched = self.TRK_Match(self.target, remainPoint, rGate, vGate, aGate)
            self.target = remainPoint[~bMatched]
            confirmed = remainPoint[bMatched]
            self.TRK_CreateNewObjects(confirmed)
        ################### Output Step ###################
        opTrace = []
        stat = self.obj['status']
        activeIdx = np.argwhere(~np.isin(stat, STAT['terminate'])).flatten()
        for i in activeIdx:
            self.trace[i].append(np.copy(self.obj[i]))
        outIdx = np.argwhere((np.isin(stat, [STAT['confirmed'], STAT['new']]) & (self.obj["sumR"]>=self.cfg["MIN_MOVED_DIST"]))).flatten()
        if len(outIdx) == 0:
            return np.zeros(0, dtype=trkOp_t), opTrace
        trkObj = np.zeros(len(outIdx), dtype=trkOp_t)
        trkObj["id"] = self.obj['id'][outIdx]
        trkObj["sumR"] = self.obj['sumR'][outIdx]
        trkObj["trkCnt"] = self.obj['detect_cnt'][outIdx] + self.obj['lost_cnt'][outIdx]
        for _n in ['distX', 'distY', 'vRelX', 'vRelY', 'snr', 'id', 'sumR']:
            trkObj[_n] = self.obj[_n][outIdx]
        if len(trkObj) == 0:
            trkObj = None
        if len(opTrace) == 0:
            opTrace = None
        return trkObj, opTrace

    def reset(self):
        self.__init__()

    def setup_cfg(self, match_dist=None, start_frame=None, min_movedist=None):
        if match_dist is not None: self.cfg["MATCH_DIST"] = match_dist
        if start_frame is not None: self.cfg["START_FRAME"] = start_frame 
        if min_movedist is not None: self.cfg["MIN_MOVED_DIST"] = min_movedist 
    
    def get_cfg(self):
        return self.cfg.copy()
